//
//  wifi.c
//
//  Created by Alexander Polyakov on 11/10/24
//
//  Copyright (c) 2024 Trident IOT
//
//
#include <stdio.h>
#include <string.h>
#include <sys/unistd.h>
#include <sys/stat.h>
#include <dirent.h>

#include "ZMalloc.h"
#include "ZWayLib.h"
#include "ZByteOperationsPrivate.h"
#include "ZFileOperationsPrivate.h"
#include "ZWayEsp32WifiPrivate.h"
#include "ZByteOperationsPrivate.h"

#include "esp_err.h"
#include "esp_log.h"
#include "esp_vfs.h"
#include "esp_spiffs.h"
#include "esp_http_server.h"
#include "esp_log.h"
#include "esp_mac.h"

#define ZWAY_ESP32_DEFAULT_AP_WIFI_CHANNEL 1
#define ZWAY_ESP32_DEFAULT_AP_MAX_CONN 5
#ifndef ZWAY_ESP32_WIFI_CONN_MAX_RETRY
#define ZWAY_ESP32_WIFI_CONN_MAX_RETRY 5
#endif
#ifndef ZWAY_ESP32_WIFI_STA_FAIL_DELAY
#define ZWAY_ESP32_WIFI_STA_FAIL_DELAY 500
#endif

#define ZWAY_ESP32_NETIF_DESC_STA "zway_netif_sta"

#define ZWAY_ESP32_WIFI_SCAN_METHOD WIFI_ALL_CHANNEL_SCAN
#define ZWAY_ESP32_WIFI_CONNECT_AP_SORT_METHOD WIFI_CONNECT_AP_BY_SIGNAL
#define ZWAY_ESP32_WIFI_SCAN_AUTH_MODE_THRESHOLD WIFI_AUTH_WPA_PSK
#define ZWAY_ESP32_WIFI_SCAN_RSSI_THRESHOLD -127

#define ZWAY_ESP32_CURRENT_WIFI_CREDENTIALS_VERSION 0x00000100

ZWError _zway_esp32_save_wifi_credentials(const ZWCSTR filename, _ZwayEsp32WifiCredentials_t *const cred)
{
    ZW_FILE_HANDLE fd;
    ZWError err = _zfile_open(NULL, &fd, filename, ZFILE_OPEN_FLAGS_WRONLY | ZFILE_OPEN_FLAGS_CREAT | ZFILE_OPEN_FLAGS_TRUNC, ZFILE_OPEN_MODE_USER_R | ZFILE_OPEN_MODE_USER_W); // keep it r/w for user, but not the group or others
    if (err != NoError)
    {
        return err;
    }
    cred->version = ZWAY_ESP32_CURRENT_WIFI_CREDENTIALS_VERSION;
    cred->crc16 = _zbyte_crc16_aug_ccitt((const ZWBYTE *)&cred->info, sizeof(cred->info));
    err = _zfile_write(NULL, fd, cred, sizeof(_ZwayEsp32WifiCredentials_t));
    _zfile_close(fd);
    return err;
}

ZWError _zway_esp32_load_wifi_credentials(const ZWCSTR filename, _ZwayEsp32WifiCredentials_t *const cred)
{
    ZW_FILE_HANDLE fd;
    ZWError err = _zfile_open(NULL, &fd, filename, ZFILE_OPEN_FLAGS_RDONLY, 0);
    if (err != NoError)
    {
        return err;
    }
    err = _zfile_read(NULL, fd, cred, sizeof(_ZwayEsp32WifiCredentials_t));
    _zfile_close(fd);
    if (err != NoError)
    {
        return err;
    }
    if (cred->version != ZWAY_ESP32_CURRENT_WIFI_CREDENTIALS_VERSION)
    {
        return BadData;
    }
    if (cred->crc16 != _zbyte_crc16_aug_ccitt((const ZWBYTE *)&cred->info, sizeof(cred->info)))
    {
        return BadData;
    }
    return NoError;
}

static void __zlog_write_information(_ZwayEsp32WifiCtx_t *ctx, ZWCSTR message)
{
    zlog_write(ctx->logger, ctx->source, Information, message);
}

#define WIFI_INIT_CONFIG_DEFAULT_DYNAMIC(cfg)                             \
    cfg->osi_funcs = &g_wifi_osi_funcs;                                   \
    cfg->wpa_crypto_funcs = g_wifi_default_wpa_crypto_funcs;              \
    cfg->static_rx_buf_num = CONFIG_ESP_WIFI_STATIC_RX_BUFFER_NUM;        \
    cfg->dynamic_rx_buf_num = CONFIG_ESP_WIFI_DYNAMIC_RX_BUFFER_NUM;      \
    cfg->tx_buf_type = CONFIG_ESP_WIFI_TX_BUFFER_TYPE;                    \
    cfg->static_tx_buf_num = WIFI_STATIC_TX_BUFFER_NUM;                   \
    cfg->dynamic_tx_buf_num = WIFI_DYNAMIC_TX_BUFFER_NUM;                 \
    cfg->rx_mgmt_buf_type = CONFIG_ESP_WIFI_DYNAMIC_RX_MGMT_BUF;          \
    cfg->rx_mgmt_buf_num = WIFI_RX_MGMT_BUF_NUM_DEF;                      \
    cfg->cache_tx_buf_num = WIFI_CACHE_TX_BUFFER_NUM;                     \
    cfg->csi_enable = WIFI_CSI_ENABLED;                                   \
    cfg->ampdu_rx_enable = WIFI_AMPDU_RX_ENABLED;                         \
    cfg->ampdu_tx_enable = WIFI_AMPDU_TX_ENABLED;                         \
    cfg->amsdu_tx_enable = WIFI_AMSDU_TX_ENABLED;                         \
    cfg->nvs_enable = WIFI_NVS_ENABLED;                                   \
    cfg->nano_enable = WIFI_NANO_FORMAT_ENABLED;                          \
    cfg->rx_ba_win = WIFI_DEFAULT_RX_BA_WIN;                              \
    cfg->wifi_task_core_id = WIFI_TASK_CORE_ID;                           \
    cfg->beacon_max_len = WIFI_SOFTAP_BEACON_MAX_LEN;                     \
    cfg->mgmt_sbuf_num = WIFI_MGMT_SBUF_NUM;                              \
    cfg->feature_caps = WIFI_FEATURE_CAPS;                                \
    cfg->sta_disconnected_pm = WIFI_STA_DISCONNECTED_PM_ENABLED;          \
    cfg->espnow_max_encrypt_num = CONFIG_ESP_WIFI_ESPNOW_MAX_ENCRYPT_NUM; \
    cfg->tx_hetb_queue_num = WIFI_TX_HETB_QUEUE_NUM;                      \
    cfg->dump_hesigb_enable = WIFI_DUMP_HESIGB_ENABLED;                   \
    cfg->magic = WIFI_INIT_CONFIG_MAGIC;

static void __wifi_init_config_init_default(wifi_init_config_t *const cfg)
{
    (void)memset(cfg, 0, sizeof(*cfg));
    WIFI_INIT_CONFIG_DEFAULT_DYNAMIC(cfg);
}

typedef struct ZwayWifiCommonStart_s
{
    union
    {
        wifi_init_config_t cfg;
        esp_netif_inherent_config_t esp_netif_config;
    };
} ZwayWifiCommonStart_t;

typedef struct ZwayWifiCommonSoftApp_s
{
    union
    {
        wifi_init_config_t cfg;
        wifi_config_t wifi_config;
    };
    ZwayEsp32Hostname_t hostname;
} ZwayWifiCommonSoftApp_t;

typedef union ZwayWifiCommon_s
{
    ZwayWifiCommonStart_t start;
    ZwayWifiCommonSoftApp_t soft_app;
} ZwayWifiCommon_t;

typedef struct ZwayWifiConnect_s
{
    _ZwayEsp32WifiCredentials_t cred;
    wifi_config_t wifi_config;
} ZwayWifiConnect_t;

typedef struct ZwayWifi_s
{
    ZwayWifiConnect_t connect;
    ZwayWifiCommon_t common;
} ZwayWifi_t;

void __wifi_sta_do_disconnect(_ZwayEsp32WifiCtx_t *ctx);

static void __handler_on_wifi_disconnect(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
    wifi_event_sta_disconnected_t *disconn;
    _ZwayEsp32WifiCtx_t *ctx;

    ctx = (_ZwayEsp32WifiCtx_t *)arg;
    disconn = event_data;
    if (disconn->reason == WIFI_REASON_ROAMING)
    {
        // station roaming, do nothing
        return;
    }
    ctx->s_retry_num++;
    if (ctx->s_semph_wait_wifi != NULL)
    {
        if (ctx->s_retry_num > ZWAY_ESP32_WIFI_CONN_MAX_RETRY)
        {
            zlog_write(ctx->logger, ctx->source, Critical, "WiFi Connect failed %d times, stop reconnect.", ctx->s_retry_num);
            __wifi_sta_do_disconnect(ctx);
            xSemaphoreGive(ctx->s_semph_wait_wifi);
            return;
        }
    }

    sleep_ms(ZWAY_ESP32_WIFI_STA_FAIL_DELAY);
    zlog_write(ctx->logger, ctx->source, Warning, "Wi-Fi disconnected %d, trying to reconnect...%d", disconn->reason, ctx->s_retry_num);
    (void)esp_wifi_connect();
    (void)event_id;
    (void)event_base;
}

static void __handler_on_sta_got_ip(void *const arg, const esp_event_base_t event_base, const int32_t event_id, void *const event_data)
{
    const ip_event_got_ip_t *const event = (ip_event_got_ip_t *)event_data;
    if (strncmp(ZWAY_ESP32_NETIF_DESC_STA, esp_netif_get_desc(event->esp_netif), (sizeof(ZWAY_ESP32_NETIF_DESC_STA) - 1)) != 0)
    {
        return;
    }
    _ZwayEsp32WifiCtx_t *const ctx = (_ZwayEsp32WifiCtx_t *)arg;
    ctx->s_retry_num = 0;
    ctx->ip4 = event->ip_info.ip;
    const SemaphoreHandle_t s_semph_wait_wifi = ctx->s_semph_wait_wifi;
    if (s_semph_wait_wifi == NULL)
    {
        return;
    }
    xSemaphoreGive(s_semph_wait_wifi);
    (void)event_id;
    (void)event_base;
}

static ZWError __wifi_start(_ZwayEsp32WifiCtx_t *ctx, ZwayWifi_t *ctx_wifi_connect)
{
    static const esp_netif_inherent_config_t esp_netif_config_default = ESP_NETIF_INHERENT_DEFAULT_WIFI_STA();
    __wifi_init_config_init_default(&ctx_wifi_connect->common.start.cfg);
    if (esp_wifi_init(&ctx_wifi_connect->common.start.cfg) != ESP_OK)
    {
        return InternalError;
    }
    
    ctx_wifi_connect->common.start.esp_netif_config = esp_netif_config_default;
    // Warning: the interface desc is used in tests to capture actual connection details (IP, gw, mask)
    ctx_wifi_connect->common.start.esp_netif_config.if_desc = ZWAY_ESP32_NETIF_DESC_STA;
    ctx_wifi_connect->common.start.esp_netif_config.route_prio = 128;
    
    if ((ctx->s_example_sta_netif = esp_netif_create_wifi(WIFI_IF_STA, &ctx_wifi_connect->common.start.esp_netif_config)) == NULL)
    {
        return InternalError;
    }
    if (esp_wifi_set_default_wifi_sta_handlers() != ESP_OK)
    {
        return InternalError;
    }

    if (esp_wifi_set_storage(WIFI_STORAGE_RAM) != ESP_OK)
    {
        return InternalError;
    }
    if (esp_wifi_set_mode(WIFI_MODE_STA) != ESP_OK)
    {
        return InternalError;
    }
    if (esp_wifi_start() != ESP_OK)
    {
        return InternalError;
    }
    return NoError;
}

static ZWError __wifi_sta_do_connect(_ZwayEsp32WifiCtx_t *ctx, ZwayWifi_t *ctx_wifi_connect)
{
    if (esp_event_handler_register(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED, &__handler_on_wifi_disconnect, ctx) != ESP_OK)
    {
        return BadAllocation;
    }
    if (esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &__handler_on_sta_got_ip, ctx) != ESP_OK)
    {
        return BadAllocation;
    }

    (void)memset(&ctx_wifi_connect->connect.wifi_config, 0, sizeof(ctx_wifi_connect->connect.wifi_config));
    ctx_wifi_connect->connect.wifi_config.sta.scan_method = ZWAY_ESP32_WIFI_SCAN_METHOD;
    ctx_wifi_connect->connect.wifi_config.sta.sort_method = ZWAY_ESP32_WIFI_CONNECT_AP_SORT_METHOD;
    ctx_wifi_connect->connect.wifi_config.sta.threshold.rssi = ZWAY_ESP32_WIFI_SCAN_RSSI_THRESHOLD;
    ctx_wifi_connect->connect.wifi_config.sta.threshold.authmode = ZWAY_ESP32_WIFI_SCAN_AUTH_MODE_THRESHOLD;
    (void)memcpy(ctx_wifi_connect->connect.wifi_config.sta.ssid, ctx_wifi_connect->connect.cred.info.ssid, strlen(ctx_wifi_connect->connect.cred.info.ssid));
    (void)memcpy(ctx_wifi_connect->connect.wifi_config.sta.password, ctx_wifi_connect->connect.cred.info.psw, strlen(ctx_wifi_connect->connect.cred.info.psw));
    static_assert((sizeof(ctx_wifi_connect->connect.wifi_config.sta.ssid)) >= (sizeof(ctx_wifi_connect->connect.cred.info.ssid)), "Overflow due to structure sizes!!!");
    static_assert((sizeof(ctx_wifi_connect->connect.wifi_config.sta.password)) >= (sizeof(ctx_wifi_connect->connect.cred.info.psw)), "Overflow due to structure sizes!!!");
    if (esp_wifi_set_config(WIFI_IF_STA, &ctx_wifi_connect->connect.wifi_config) != ESP_OK)
    {
        return InternalError;
    }
    if (esp_wifi_connect() != ESP_OK)
    {
        return InternalError;
    }
    __zlog_write_information(ctx, "Waiting for IP(s)...");
    xSemaphoreTake(ctx->s_semph_wait_wifi, portMAX_DELAY);
    if (ctx->s_retry_num > ZWAY_ESP32_WIFI_CONN_MAX_RETRY)
    {
        return InternalError;
    }
    vSemaphoreDelete(ctx->s_semph_wait_wifi);
    ctx->s_semph_wait_wifi = NULL;
    return NoError;
}

void __wifi_sta_do_disconnect(_ZwayEsp32WifiCtx_t *ctx)
{
    (void)(esp_event_handler_unregister(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED, &__handler_on_wifi_disconnect));
    (void)(esp_event_handler_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, &__handler_on_sta_got_ip));
    (void)esp_wifi_disconnect();
}

static void __wifi_shutdown(_ZwayEsp32WifiCtx_t *ctx)
{
    SemaphoreHandle_t s_semph_wait_wifi;
    esp_netif_t *s_example_sta_netif;

    (void)esp_wifi_stop();
    (void)esp_wifi_deinit();
    if ((s_example_sta_netif = ctx->s_example_sta_netif) != NULL)
    {
        (void)esp_wifi_clear_default_wifi_driver_and_handlers(s_example_sta_netif);
        esp_netif_destroy(s_example_sta_netif);
    }
    if ((s_semph_wait_wifi = ctx->s_semph_wait_wifi) != NULL)
    {
        vSemaphoreDelete(s_semph_wait_wifi);
    }
}

static void __zway_esp32_wifi_shutdown(_ZwayEsp32WifiCtx_t *ctx)
{
    if (ctx->inited)
    {
        __wifi_sta_do_disconnect(ctx);
        __wifi_shutdown(ctx);
        ctx->inited = false;
    }
    zfree(ctx);
}

static ZWError __wifi_init_softap(_ZwayEsp32WifiCtx_t *ctx, ZwayWifi_t *ctx_wifi_connect)
{
    ZWError err;

    if ((err = _zway_esp32_get_hostname(&ctx_wifi_connect->common.soft_app.hostname)) != NoError)
    {
        return err;
    }

    if ((ctx->s_example_sta_netif = esp_netif_create_default_wifi_ap()) == NULL)
    {
        return BadAllocation;
    }
    __wifi_init_config_init_default(&ctx_wifi_connect->common.soft_app.cfg);
    if (esp_wifi_init(&ctx_wifi_connect->common.soft_app.cfg) != ESP_OK)
    {
        return InternalError;
    }

    static_assert((sizeof(ctx_wifi_connect->common.soft_app.wifi_config.ap.ssid)) >= (sizeof(ctx_wifi_connect->common.soft_app.hostname)), "Overflow due to structure sizes!!!");
    static_assert((sizeof(ctx_wifi_connect->common.soft_app.wifi_config.ap.password)) >= (sizeof(ZWAY_ESP32_DEFAULT_AP_WIFI_PASS)), "Overflow due to structure sizes!!!");
    (void)memset(&ctx_wifi_connect->common.soft_app.wifi_config, 0, sizeof(ctx_wifi_connect->common.soft_app.wifi_config));
    ctx_wifi_connect->common.soft_app.wifi_config.ap.channel = ZWAY_ESP32_DEFAULT_AP_WIFI_CHANNEL;
    (void)memcpy(ctx_wifi_connect->common.soft_app.wifi_config.ap.password, ZWAY_ESP32_DEFAULT_AP_WIFI_PASS, sizeof(ZWAY_ESP32_DEFAULT_AP_WIFI_PASS));
    ctx_wifi_connect->common.soft_app.wifi_config.ap.max_connection = ZWAY_ESP32_DEFAULT_AP_MAX_CONN;
    ctx_wifi_connect->common.soft_app.wifi_config.ap.authmode = WIFI_AUTH_WPA2_PSK;
    ctx_wifi_connect->common.soft_app.wifi_config.ap.pmf_cfg.required = true;
    ctx_wifi_connect->common.soft_app.wifi_config.ap.ssid_len = strlen(ctx_wifi_connect->common.soft_app.hostname.data);
    (void)memcpy(ctx_wifi_connect->common.soft_app.wifi_config.ap.ssid, ctx_wifi_connect->common.soft_app.hostname.data, ctx_wifi_connect->common.soft_app.wifi_config.ap.ssid_len);
    if (sizeof(ZWAY_ESP32_DEFAULT_AP_WIFI_PASS) == 1)
    {
        ctx_wifi_connect->common.soft_app.wifi_config.ap.authmode = WIFI_AUTH_OPEN;
    }

    if (esp_wifi_set_mode(WIFI_MODE_AP) != ESP_OK)
    {
        return InternalError;
    }
    if (esp_wifi_set_config(WIFI_IF_AP, &ctx_wifi_connect->common.soft_app.wifi_config) != ESP_OK)
    {
        return InternalError;
    }
    if (esp_wifi_start() != ESP_OK)
    {
        return InternalError;
    }
    return NoError;
}

static ZWError __zway_esp32_wifi_connect(_ZwayEsp32WifiCtx_t *ctx, ZWBOOL force_ap, ZwayWifi_t *ctx_wifi_connect)
{
   
    if (force_ap)
    {
        return __wifi_init_softap(ctx, ctx_wifi_connect);
    }
    if (_zway_esp32_load_wifi_credentials(ZWAY_ESP32_WIFI_CREDENTIAL_FILENAME, &ctx_wifi_connect->connect.cred) != NoError)
    {
        return BadData;
    }
    if ((ctx->s_semph_wait_wifi = xSemaphoreCreateBinary()) == NULL)
    {
        return BadAllocation;
    }
    ZWError err = __wifi_start(ctx, ctx_wifi_connect);
    if (err != NoError)
    {
        return err;
    }
    ctx->inited = true;
    if ((err = __wifi_sta_do_connect(ctx, ctx_wifi_connect)) != NoError)
    {
        return err;
    }
    return (NoError);
}

ZWError _zway_esp32_wifi_connect(_ZwayEsp32WifiCtx_t **ctx, ZWLog logger, ZWCSTR source, ZWBOOL force_ap)
{
    _ZwayEsp32WifiCtx_t *ctx_local;
    ZwayWifi_t *ctx_wifi_connect;
    ZWError err;

    if ((ctx_local = zmalloc(sizeof(_ZwayEsp32WifiCtx_t))) == NULL)
    {
        return (BadAllocation);
    }
    if ((ctx_wifi_connect = zmalloc(sizeof(ZwayWifi_t))) == NULL)
    {
        zfree(ctx_local);
        return (BadAllocation);
    }
    (void)memset(ctx_local, 0, sizeof(_ZwayEsp32WifiCtx_t));
    ctx_local->logger = logger;
    ctx_local->source = source;
    if ((err = __zway_esp32_wifi_connect(ctx_local, force_ap, ctx_wifi_connect) != NoError))
    {
        __zway_esp32_wifi_shutdown(ctx_local);
    }
    else
    {
        *ctx = ctx_local;
    }
    zfree(ctx_wifi_connect);
    return (err);
}

ZWError _zway_esp32_get_hostname(ZwayEsp32Hostname_t *hostname)
{
    ZWBYTE mac[8];

    if (esp_efuse_mac_get_default(mac) != ESP_OK)
    {
        return InternalError;
    }
    if (snprintf(hostname->data, sizeof(hostname->data), "ZWAY-ESP-%02x%02x", mac[4], mac[5]) != 13)
    {
        return InvalidArg;
    }
    return NoError;
}